{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "# In Google Colab, uncomment this:\n",
    "# !wget https://bit.ly/2FMJP5K -O setup.py && bash setup.py\n",
    "\n",
    "# This code creates a virtual display to draw game images on.\n",
    "# If you are running locally, just ignore it\n",
    "import os\n",
    "if type(os.environ.get(\"DISPLAY\")) is not str or len(os.environ.get(\"DISPLAY\")) == 0:\n",
    "    !bash ../xvfb start\n",
    "    os.environ['DISPLAY'] = ':1'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OpenAI Gym\n",
    "\n",
    "We're gonna spend several next weeks learning algorithms that solve decision processes. We are then in need of some interesting decision problems to test our algorithms.\n",
    "\n",
    "That's where OpenAI gym comes into play. It's a python library that wraps many classical decision problems including robot control, videogames and board games.\n",
    "\n",
    "So here's how it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "\n",
    "env = gym.make(\"MountainCar-v0\")\n",
    "env.reset()\n",
    "\n",
    "plt.imshow(env.render('rgb_array'))\n",
    "print(\"Observation space:\", env.observation_space)\n",
    "print(\"Action space:\", env.action_space)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: if you're running this on your local machine, you'll see a window pop up with the image above. Don't close it, just alt-tab away."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gym interface\n",
    "\n",
    "The three main methods of an environment are\n",
    "* __reset()__ - reset environment to initial state, _return first observation_\n",
    "* __render()__ - show current environment state (a more colorful version :) )\n",
    "* __step(a)__ - commit action __a__ and return (new observation, reward, is done, info)\n",
    " * _new observation_ - an observation right after commiting the action __a__\n",
    " * _reward_ - a number representing your reward for commiting action __a__\n",
    " * _is done_ - True if the MDP has just finished, False if still in progress\n",
    " * _info_ - some auxilary stuff about what just happened. Ignore it ~~for now~~."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs0 = env.reset()\n",
    "print(\"initial observation code:\", obs0)\n",
    "\n",
    "# Note: in MountainCar, observation is just two numbers: car position and velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"taking action 2 (right)\")\n",
    "new_obs, reward, is_done, _ = env.step(2)\n",
    "\n",
    "print(\"new observation code:\", new_obs)\n",
    "print(\"reward:\", reward)\n",
    "print(\"is game over?:\", is_done)\n",
    "\n",
    "# Note: as you can see, the car has moved to the right slightly (around 0.0005)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Play with it\n",
    "\n",
    "Below is the code that drives the car to the right. \n",
    "\n",
    "However, it doesn't reach the flag at the far right due to gravity. \n",
    "\n",
    "__Your task__ is to fix it. Find a strategy that reaches the flag. \n",
    "\n",
    "You're not required to build any sophisticated algorithms for now, feel free to hard-code :)\n",
    "\n",
    "__Hint__: your action at each step should depend either on `t` or on `s`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython import display\n",
    "\n",
    "# create env manually to set time limit. Please don't change this.\n",
    "TIME_LIMIT = 250\n",
    "env = gym.wrappers.TimeLimit(\n",
    "    gym.envs.classic_control.MountainCarEnv(),\n",
    "    max_episode_steps=TIME_LIMIT + 1,\n",
    ")\n",
    "s = env.reset()\n",
    "actions = {'left': 0, 'stop': 1, 'right': 2}\n",
    "\n",
    "plt.figure(figsize=(4, 3))\n",
    "display.clear_output(wait=True)\n",
    "\n",
    "for t in range(TIME_LIMIT):\n",
    "    plt.gca().clear()\n",
    "    \n",
    "    # change the line below to reach the flag\n",
    "    s, r, done, _ = env.step(actions['right'])\n",
    "\n",
    "    # draw game image on display\n",
    "    plt.imshow(env.render('rgb_array'))\n",
    "    \n",
    "    display.clear_output(wait=True)\n",
    "    display.display(plt.gcf())\n",
    "\n",
    "    if done:\n",
    "        print(\"Well done!\")\n",
    "        break\n",
    "else:\n",
    "    print(\"Time limit exceeded. Try again.\")\n",
    "\n",
    "display.clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert s[0] > 0.47\n",
    "print(\"You solved it!\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
